import asyncio
import logging
from typing import Dict, Any
from openai import OpenAI
from video_agent.config import (
    MODEL_NAME, MODEL_TIMEOUT, API_BASE_URL, API_KEY
)

logger = logging.getLogger(__name__)

class ModelCaller:
    def __init__(self):
        self.client = OpenAI(
            api_key=API_KEY,
            base_url=API_BASE_URL
        )
    
    async def create_completion(
        self, 
        video_base64: str, 
        text_prompt: str
    ) -> Dict[str, Any]:
        """创建模型调用"""
        try:
            # 创建流式响应
            response = self.client.chat.completions.create(
                model=MODEL_NAME,
                messages=[
                    {
                        "role": "system",
                        "content": [{"type": "text", "text": "You are a helpful assistant."}],
                    },
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "video_url",
                                "video_url": {"url": f"data:;base64,{video_base64}"},
                            },
                            {"type": "text", "text": text_prompt},
                        ],
                    },
                ],
                modalities=["text"],
                stream=True,  # 必须设置为 True
                stream_options={"include_usage": True},
            )
            
            # 处理流式响应
            text_response = ""
            usage = None
            
            # 使用超时控制
            async with asyncio.timeout(MODEL_TIMEOUT):
                for chunk in response:  # 使用普通的 for 循环，因为 response 是同步迭代器
                    if chunk.choices:
                        delta = chunk.choices[0].delta
                        if hasattr(delta, 'content') and delta.content:
                            text_response += delta.content
                    else:
                        usage = chunk.usage
            
            if not text_response:
                logger.error("没有收到有效的文本响应")
                raise Exception("模型返回的响应为空")
            
            logger.info(f"最终响应文本: {text_response}")
            return {
                "text": text_response,
                "usage": usage
            }
            
        except asyncio.TimeoutError:
            logger.error(f"模型调用超时: {MODEL_TIMEOUT}秒")
            raise
        except Exception as e:
            logger.error(f"模型调用失败: {str(e)}")
            raise 